#include "Containers/RangePair.h"
#include "File/BOSSOutputLoader.h"
#include "Fit/RooFitBuilder.h"
#include "Functions/Draw.h"
#include "Functions/Error.h"
#include "Functions/TerminalIO.h"
#include "Globals/Settings.h"
#include "Particle/ReconstructedParticle.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TStyle.h"
#include <iostream>

using namespace BOSS_Afterburner;
using namespace BOSS_Afterburner::Draw;
using namespace BOSS_Afterburner::Error;
using namespace BOSS_Afterburner::TerminalIO;
using namespace RooFit;
using namespace std;

/// @addtogroup Afterburner_scripts
/// @{
/// @author   Remco de Boer 雷穆克 (r.e.deboer@students.uu.nl or remco.de.boer@ihep.ac.cn)
/// @date     May 29th, 2019
void Model(const char* config_file)
{
  YamlLoader       conf(config_file);
  BOSSOutputLoader file(conf);

  // * Get particle and draw options * //
  ReconstructedParticle particle(conf["Reconstructed particle"]);
  BranchPlotOptions     drawOpts(conf["Plot to analyse"]);

  auto        signalHist = file.Draw(drawOpts);
  const char* unit       = drawOpts[0].GetUnit();

  RangePair fit  = conf["Ranges"]["fit"];
  RangePair plot = conf["Ranges"]["plot"];

  // * Construct data set + variable * //
  const char* varname =
    Form("#it{M}_{%s #rightarrow %s}", particle.NameLaTeX(), particle.DaughterLabel());
  RooRealVar    var(varname, varname, drawOpts[0].From(), drawOpts[0].To(), unit);
  RooFitBuilder builder(var);

  // * Construct RooFit data histogram * //
  const auto& dataDef = conf["Plot to analyse"];
  if(!dataDef || !dataDef.IsMap()) EXCEPTION("Missing key \"Plot to analyse\"");
  auto        hist = file.Draw(dataDef);
  RooDataHist dataHist("data", hist->GetTitle(), var, Import(*hist));

  // * Build complete fit model * //
  const auto& modelDef = conf["Fit model"];
  if(!modelDef || !modelDef.IsMap()) EXCEPTION("Missing key \"Fit model\"");
  SigPlusBck model = builder.BuildSigPlusBck(modelDef);

  // * Plot data * //
  RooPlot* frame = var.frame();
  dataHist.plotOn(frame, LineWidth(2), LineColor(kBlue + 2), LineWidth(1), MarkerColor(kBlue + 2),
                  MarkerSize(.5));
  PlotStatsOnFrame(dataHist, frame);

  // * Legend position * //
  Double_t xMin     = .68;
  Double_t xMax     = .97;
  Double_t yMin     = .92;
  Double_t textsize = .033;

  // * Perform fit convoluted with truncated polynomial background * //
  Color_t colour;
  colour = kRed;
  model.pdf->fitTo(dataHist, Range(fit.X1(), fit.X2()));
  model.pdf->plotOn(frame, LineWidth(2), LineColor(colour), Range(plot.X1(), plot.X2()),
                    LineStyle(kDotted));
  model.pdf->plotOn(frame, LineWidth(2), LineColor(colour));
  if(model.bck.pdf)
    model.pdf->plotOn(frame, Components(*model.bck.pdf), LineStyle(kDotted), LineWidth(2),
                      LineColor(kBlack));

  model.pdf->paramOn(frame, Layout(xMin, xMax, yMin), Format("NU", AutoPrecision(3)));
  if(frame)
  {
    frame->getAttText()->SetTextSize(textsize);
    frame->getAttText()->SetTextColor(colour);
  }

  TCanvas c;
  c.SetBatch();
  frame->Draw();

  // * Draw indications * //
  const auto& srDef = conf["3-sigma range"];
  try
  {
    Indicate3SigmaRange(conf);
  }
  catch(const YAML::TypedBadConversion<Double_t>& e)
  {
    RooArgSet data;
    auto      params = model.pdf->getParameters(data);
    double    width  = 0.;
    for(const auto& sig : srDef["sigmas"])
    {
      auto par = dynamic_cast<RooRealVar*>(params->find(sig.as<string>().c_str()));
      if(par) width += par->getVal();
    }
    auto par = dynamic_cast<RooRealVar*>(params->find(srDef["mass"].as<string>().c_str()));
    if(!par) throw YAML::BadConversion(YAML::Mark::null_mark());
    Indicate3SigmaRange(par->getVal(), width);
  }
  IndicateFitRange(conf);
  IndicateMassWindow(conf);

  // * Draw and write fitted histograms * //
  TString outputName;
  if(model.sig.pdf) outputName += model.sig.name + "_";
  if(model.bck.pdf) outputName += model.bck.name + "_";
  outputName.ReplaceAll(" ", "");
  SaveCanvas(Form("%s%s", outputName.Data(), particle.Name()), &c, drawOpts.LogXYZ().c_str());
}
/// @}

int main(int argc, char* argv[])
{
  CheckMainArguments(argc, argv, 1);
  Model(argv[1]);
  return 0;
}